import threading
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional


def _timestamp() -> str:
    """Return an ISO8601 UTC timestamp string."""
    return datetime.now(timezone.utc).isoformat()


class ProgressTracker:
    """Thread-safe tracker for node and tool progress events."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._sessions: Dict[str, Dict[str, Any]] = {}

    def start(self, job_id: Optional[str], metadata: Optional[Dict[str, Any]] = None) -> None:
        if not job_id:
            return
        with self._lock:
            self._sessions[job_id] = {
                "events": [],
                "finished": False,
                "seq": 0,
                "metadata": metadata or {},
            }
            self._append_event_locked(job_id, {"type": "status", "status": "started"})

    def _append_event_locked(self, job_id: str, payload: Dict[str, Any]) -> None:
        session = self._sessions.setdefault(
            job_id,
            {"events": [], "finished": False, "seq": 0, "metadata": {}},
        )
        session["seq"] += 1
        event = {
            "seq": session["seq"],
            "timestamp": _timestamp(),
            **payload,
        }
        session["events"].append(event)

    def _append_event(self, job_id: Optional[str], payload: Dict[str, Any]) -> None:
        if not job_id:
            return
        with self._lock:
            self._append_event_locked(job_id, payload)

    def record_node(self, job_id: Optional[str], node_name: str) -> None:
        self._append_event(job_id, {"type": "node", "node": node_name})

    def record_tool(
        self,
        job_id: Optional[str],
        tool_name: str,
        *,
        status: str = "success",
        arguments: Optional[Any] = None,
        detail: Optional[str] = None,
    ) -> None:
        payload: Dict[str, Any] = {"type": "tool", "tool": tool_name, "status": status}
        if arguments is not None:
            payload["arguments"] = arguments
        if detail:
            payload["detail"] = detail
        self._append_event(job_id, payload)

    def record_assistance_request(
        self,
        job_id: Optional[str],
        request_id: str,
        prompt: str,
        context: Optional[str] = None,
    ) -> None:
        payload: Dict[str, Any] = {
            "type": "human_assist_request",
            "requestId": request_id,
            "prompt": prompt,
        }
        if context:
            payload["context"] = context
        self._append_event(job_id, payload)

    def record_message(self, job_id: Optional[str], message: str) -> None:
        self._append_event(job_id, {"type": "info", "message": message})

    def finish(self, job_id: Optional[str], status: str = "success", message: Optional[str] = None) -> None:
        if not job_id:
            return
        with self._lock:
            self._append_event_locked(
                job_id,
                {
                    "type": "status",
                    "status": status,
                    "message": message,
                },
            )
            session = self._sessions.get(job_id)
            if session:
                session["finished"] = True

    def fail(self, job_id: Optional[str], message: Optional[str] = None) -> None:
        self.finish(job_id, status="failed", message=message)

    def get_events(self, job_id: str, after_seq: int = 0) -> Dict[str, Any]:
        with self._lock:
            session = self._sessions.get(job_id)
            if not session:
                return {"events": [], "finished": True, "missing": True}
            events = [evt for evt in session["events"] if evt["seq"] > after_seq]
            latest_seq = session["seq"]
            result = {
                "events": events,
                "finished": session["finished"],
                "latestSeq": latest_seq,
                "metadata": session.get("metadata", {}),
            }
            return result


progress_tracker = ProgressTracker()
